package com.udf.flink.udaf;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.table.annotation.DataTypeHint;
import org.apache.flink.table.annotation.FunctionHint;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.api.EnvironmentSettings;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.TableEnvironment;
import org.apache.flink.table.functions.AggregateFunction;
import org.apache.flink.table.functions.ScalarFunction;
import org.apache.flink.table.functions.TableFunction;
import org.apache.flink.table.functions.UserDefinedFunction;
import org.apache.flink.types.Row;

import java.time.LocalDate;


/**
 * Example for implementing more complex {@link UserDefinedFunction}s.
 *
 * <p>In many use cases, function signatures can be reflectively extracted from a UDF class. The
 * annotations {@link DataTypeHint} and {@link FunctionHint} help if reflective information is not
 * enough and needs to be enriched with further logical details. Check the website documentation as
 * well as the docs of {@link ScalarFunction}, {@link TableFunction}, and {@link AggregateFunction}
 * for more information.
 *
 * <p>Both reflective extraction and annotations are suitable for function signatures with fixed
 * input and output types. However, for advanced use cases it might be required to derive an output
 * type from one of the argument types or perform stricter validation.
 *
 * <p>This example demonstrates various UDF implementations. We are executing multiple Flink jobs
 * where the result is written to stdout.
 */
public class AdvancedFunctionsExample {

    public static void main(String[] args) throws Exception {
        // setup the environment
        final EnvironmentSettings settings = EnvironmentSettings.newInstance()
                .inBatchMode()
                .build();
        final TableEnvironment tEnv = TableEnvironment.create(settings);
        Configuration configuration = tEnv.getConfig().getConfiguration();
        configuration.setString("table.exec.mini-batch.enabled", "true");
        configuration.setString("table.exec.mini-batch.allow-latency", "5 s");
        configuration.setString("table.exec.mini-batch.size", "5000");
        configuration.setString("sql-client.execution.result-mode","TABLEAU");

        // execute different kinds of functions
        executeLastDatedValueFunction(tEnv);
//        executeInternalRowMergerFunction(tEnv);
    }

    /**
     * Aggregates data by name and returns the latest non-null {@code item_count} value with its
     * corresponding {@code order_date}.
     */
    private static void executeLastDatedValueFunction(TableEnvironment tEnv) {
        // create a table with example data
        final Table customers =
                tEnv.fromValues(
                        DataTypes.of("ROW<name STRING, order_date DATE, item_count INT,item_price BIGINT>"),
                        Row.of("Guillermo Smith", LocalDate.parse("2020-12-01"), 3,2),
                        Row.of("Guillermo Smith", LocalDate.parse("2020-12-05"), 5,3),
                        Row.of("Valeria Mendoza", LocalDate.parse("2020-03-23"), 4,3),
                        Row.of("Valeria Mendoza", LocalDate.parse("2020-06-02"), 10,2),
                        Row.of("Leann Holloway", LocalDate.parse("2020-05-26"), 9,1),
                        Row.of("Leann Holloway", LocalDate.parse("2020-05-27"), 1,2),
                        Row.of("Brandy Sanders", LocalDate.parse("2020-10-14"), 1,3),
                        Row.of("John Turner", LocalDate.parse("2020-10-02"), 12,3),
                        Row.of("Ellen Ortega", LocalDate.parse("2020-06-18"), 100,3));
        tEnv.createTemporaryView("customers", customers);


        // register and execute the function
        tEnv.createTemporarySystemFunction("LastDatedValueFunction", LastDatedValueFunction.class);
        tEnv.createTemporarySystemFunction("weightedAvg", WeightedAvgFunction.class);
//        String sql = "SELECT name, LastDatedValueFunction(item_count, order_date) FROM customers GROUP BY name";
        String sql = "SELECT name, weightedAvg(item_price, item_count) FROM customers GROUP BY name";
        tEnv.executeSql(sql).print();

        // clean up
        tEnv.dropTemporaryView("customers");
    }

    /**
     * Merges two rows as efficient as possible using internal data structures.
     */
    private static void executeInternalRowMergerFunction(TableEnvironment env) {
        // create a table with example data
        final Table customers =
                env.fromValues(
                        DataTypes.of(
                                "ROW<name STRING, data1 ROW<birth_date DATE>, data2 ROW<city STRING, phone STRING>>"),
                        Row.of(
                                "Guillermo Smith",
                                Row.of(LocalDate.parse("1992-12-12")),
                                Row.of("New Jersey", "816-443-8010")),
                        Row.of(
                                "Valeria Mendoza",
                                Row.of(LocalDate.parse("1970-03-28")),
                                Row.of("Los Angeles", "928-264-9662")),
                        Row.of(
                                "Leann Holloway",
                                Row.of(LocalDate.parse("1989-05-21")),
                                Row.of("Eugene", "614-889-6038")));
        env.createTemporaryView("customers", customers);

        // register and execute the function
        env.createTemporarySystemFunction(
                "InternalRowMergerFunction", InternalRowMergerFunction.class);
        env.executeSql("SELECT name, InternalRowMergerFunction(data1, data2) FROM customers")
                .print();

        // clean up
        env.dropTemporaryView("customers");
    }
}
